package cn.ratel.tool;

import cn.ratel.kit.StrKit;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * Created by iver
 */
public class ClassSearcher {
	private static Logger log = LoggerFactory.getLogger(ClassSearcher.class);

	public static Set<Class<?>> scan(String packageName, Filter filter) {
		try {
			return doScan(packageName, filter);
		} catch (IOException | ClassNotFoundException e) {
			log.error("scan class error : " + e.getMessage(), e);
		}
		return null;
	}

	private static Set<Class<?>> doScan(String packageName, Filter filter) throws IOException, ClassNotFoundException {
		packageName = StrKit.defaultStr(packageName);
		if (StrKit.notBlank(packageName) && !packageName.endsWith(".")) {
			packageName += ".";
		}
		Enumeration<URL> resources = getClassLoader().getResources(packageName.replace('.', '/'));
		if (resources == null) {
			return null;
		}
		Set<Class<?>> classes = new HashSet<>();
		while (resources.hasMoreElements()) {
			String path = resources.nextElement().getPath();
			int jarIndex = path.indexOf(".jar!");
			if (jarIndex != -1) {
				String jarFileName = path.replace("file:/", "").substring(0, jarIndex - 2);
				addJarClasses(packageName, filter, classes, new JarFile(jarFileName));
			} else {
				addClasses(packageName, filter, classes, path);
			}
		}
		return classes;
	}

	private static void addJarClasses(String packageName, Filter filter, Set<Class<?>> classes, JarFile jarFile) throws IOException, ClassNotFoundException {
		Enumeration<JarEntry> entries = jarFile.entries();
		while (entries.hasMoreElements()) {
			JarEntry entry = entries.nextElement();
			String className = entry.getName().replace('/', '.');
			if (className.endsWith(".class") && className.startsWith(packageName)) {
				Class<?> clazz = getClassLoader().loadClass(className.substring(0, className.length() - 6));
				if (filter.doFilter(clazz)) {
					classes.add(clazz);
				}
			}
		}
	}

	private static void addClasses(String packageName, Filter filter, Set<Class<?>> classes, String path) throws ClassNotFoundException {
		path = path.substring(1, path.length());
		Collection<File> classFiles = FileUtils.listFiles(new File(path), new String[]{"class"}, true);
		if (classFiles != null) {
			for (File classFile : classFiles) {
				String file = classFile.getPath();
				String fileWithDot = file.replace(File.separatorChar, '.');
				String className = packageName.trim() + fileWithDot.substring(path.length(), file.length() - 6);
				Class<?> clazz = getClassLoader().loadClass(className);
				if (filter.doFilter(clazz)) {
					classes.add(clazz);
				}
			}
		}
	}

	private static ClassLoader getClassLoader() {
		ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
		if (classLoader == null) {
			classLoader = ClassSearcher.getClassLoader();
			if (null == classLoader) {
				classLoader = ClassLoader.getSystemClassLoader();
			}
		}
		return classLoader;
	}

	public interface Filter {
		boolean doFilter(Class<?> clazz);
	}

}
